#! /usr/bin/env python3

import rospy
import requests
import threading


from std_msgs.msg import String
from cozepy import COZE_CN_BASE_URL
from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageContentType, ChatEventType  # noqa

# TOKEN = "pat_BpQRn0LEWYZMmcb8teFpaCOHqp01UtOM6zHhhRBKy9G4c7r7z8RmstKXoAJuvG3Y"
# TOKEN = "pat_OKhkPUe0kZFESsYiKjIeggrA30MuXjYvawBZ163fzjHVLnC3tVKVtRUcCXjSa4Q2"
# TOKEN = "pat_84Uq0NNDhevxItN4HMxPdk0uu73L6qY4k4BXHa9HJETy5a1I8jN0QcZi9DAwfoRX"
# TOKEN = "pat_fFQAViNLD0Ochfgr91Ko9qaKmzj6NnXWs9GEYIGzoFp3UzJdRzS6SJTASIGwUlHM"

TOKEN = "sat_bLF9J7Vziicgya22Qe2k9V1Ylkd6oZkAcNesw2jLe62izArO4PJNUngownYIWZLq"

class CozeAgent():


    def __init__(self):
        #智能体id
        self.bot_id = "7485202131185123340"
        self.user_id = '001'
        self.headers = {
            "Authorization": f"Bearer {TOKEN}"
        }

        # 创建会话
        self.conversion_id = self.create_conversion_id()

        self.coze_api_token = TOKEN
        self.coze_api_base = COZE_CN_BASE_URL

        # Init the Coze client through the access_token.
        self.coze = Coze(auth = TokenAuth(token = self.coze_api_token), 
                         base_url = self.coze_api_base)

        # 初始化 ros 节点和 发布者 和 订阅者
        rospy.init_node("know_llm_node")
        rospy.loginfo("huzhou llm node start......")

        # 发布大模型的回答
        self.pub = rospy.Publisher("huzhou_llm_res", String, queue_size=10)
        
        # 发布 ASR 识别到的信息
        self.sub = rospy.Subscriber("huzhou_llm_req", String, self.receiveQuestionAndPost, queue_size=10)
        
        # 接受是否停止模型输出
        self.stop_sub = rospy.Subscriber("huzhou_llm_stop", String, self.receiveStopCommand, queue_size=10)


        self.lock = threading.Lock()
        self.stop_lock = threading.Lock()

        self.stop_flag = False  # 用于指示是否终止 LLM 回答

        self.is_running = False  # 指示是否处于大模型的访问时期


    def receiveStopCommand(self, msg):
        # 当接收到停止命令时，修改 stop_flag
        with self.stop_lock:  # 保护 stop_flag 和 is_running
            if self.is_running:
                self.stop_flag = True
                print("receive stop flag when running ......")
    

    def receiveQuestionAndPost(self, msg):
        rospy.loginfo("receive question: %s", msg.data)

        # 新开一个线程来运行代码
        threading.Thread(target=self.llmPost, args=(msg.data,), daemon=True).start()
        # self.llmPost(msg.data)
        return


    # 向大语言模型服务器发送请求
    def llmPost(self, question):
    #     msg = String()
    #     msg.data = "肾病患者就更应格外注意日常生活中的点滴。要养成良好的生活习惯，戒烟限酒，避免劳累（熬夜）。适当运动，定期监测血压，及时和医生沟通。应适当补充维生素、矿物质及微量元素，肾病综合征患者补充维生素D3和钙很重要，同时应适量补充维生素A、B、C及锌。以上是肾病患者应该注意的事项。"  
    #     self.pub.publish(msg)
    #     msg.data = "[DONE]"
    #     self.pub.publish(msg)

        with self.lock:  # 线程安全

            self.is_running = True

            for event in self.coze.chat.stream(
                    bot_id = self.bot_id,
                    user_id = self.user_id,
                    conversion_id = self.conversion_id,
                    additional_messages=[
                        Message.build_user_question_text(question),
                    ],
            ):
                with self.stop_lock:  # 保护 stop_flag 和 is_running
                    # 接受到停止的信号，应该立刻停止
                    if self.stop_flag:
                        rospy.loginfo("Stopping LLM response early.")
                        self.stop_flag = False

                        self.is_running = False

                        # 发送截至符号
                        msg = String()
                        msg.data = "[DONE]"
                        self.pub.publish(msg)

                        break

                if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:

                    # 打印出从llm服务器返回的数据
                    print(event.message.content, end="", flush=True)

                    # 将大模型的回答发送给 ros 订阅者
                    msg = String()
                    msg.data = event.message.content
                    self.pub.publish(msg)

                if event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
                    self.is_running = False

                    print()
                    print("token usage:", event.chat.usage.token_count)

                    # 发送截至符号
                    msg = String()
                    msg.data = "[DONE]"
                    self.pub.publish(msg)
        #end of 'self.lock'


    #第一步 创建会话
    def create_conversion_id(self):
        url = COZE_CN_BASE_URL + "/v1/conversation/create"

        response = requests.post(url, headers=self.headers)
        if response.status_code == 200:
            data = response.json()
            print(data)
            return data["data"]["id"]
        else:
            print(f"Error: {response.status_code}")
            return None


if __name__ == "__main__":

    coze_agent = CozeAgent()

    # coze_agent.llmPost("湖州吴兴AI+知识产权服务业集聚区建设情况");

    rospy.spin();
